import torch
import torch.nn as nn

class SimpleTransformerActor(nn.Module):
    def __init__(self, input_dim=96, hidden_dim=128, n_layers=2, n_heads=4, n_actions=12):
        super().__init__()
        self.embed = nn.Linear(input_dim, hidden_dim)
        encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=n_heads)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.output = nn.Linear(hidden_dim, n_actions)

    def forward(self, x):
        # x: [batch_size, seq_len, input_dim]
        x = self.embed(x)  # [B, T, H]
        x = x.permute(1, 0, 2)  # [T, B, H]
        x = self.transformer(x)  # [T, B, H]
        x = x.permute(1, 0, 2)  # [B, T, H]
        out = self.output(x)  # [B, T, n_actions]
        return out
